{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tracking token usage \n",
    "跟踪令牌使用情况\n",
    "> This notebook goes over how to track your token usage for specific calls. It is currently only implemented for the OpenAI API.<br>\n",
    "本笔记本介绍了如何跟踪特定调用的令牌使用情况。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Why was the math book sad? Because it had too many problems.\n",
      "---\n",
      "\n",
      "Total Tokens: 19\n",
      "Prompt Tokens: 4\n",
      "Completion Tokens: 15\n",
      "Total Cost (USD): $3.6e-05\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.callbacks import get_openai_callback\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model_name=\"gpt-3.5-turbo-instruct\")\n",
    "\n",
    "with get_openai_callback() as cb:\n",
    "    result = llm.invoke(\"Tell me a joke\")\n",
    "    print(result)\n",
    "    print(\"---\")\n",
    "print()\n",
    "\n",
    "print(f\"Total Tokens: {cb.total_tokens}\")\n",
    "print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
    "print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
    "print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Why did the chicken cross the playground?\n",
      "\n",
      "To get to the other slide!\n",
      "--\n",
      "\n",
      "\n",
      "Why did the fish get bad grades?\n",
      "\n",
      "Because it was always swimming in schools!\n",
      "\n",
      "---\n",
      "Total Tokens: 45\n",
      "Prompt Tokens: 12\n",
      "Completion Tokens: 33\n",
      "Total Cost (USD): $8.400000000000001e-05\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.callbacks import get_openai_callback\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model_name=\"gpt-3.5-turbo-instruct\")\n",
    "\n",
    "template = PromptTemplate.from_template(\"Tell me a joke about {topic}\")\n",
    "chain = template | llm\n",
    "\n",
    "with get_openai_callback() as cb:\n",
    "    response = chain.invoke({\"topic\": \"birds\"})\n",
    "    print(response)\n",
    "    response = chain.invoke({\"topic\": \"fish\"})\n",
    "    print(\"--\")\n",
    "    print(response)\n",
    "\n",
    "\n",
    "print()\n",
    "print(\"---\")\n",
    "print(f\"Total Tokens: {cb.total_tokens}\")\n",
    "print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
    "print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
    "print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming\n",
    "\n",
    "> get_openai_callback 目前不支持旧语言模型（例如 langchain_openai.OpenAI ）的流式令牌计数。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Why couldn't the bicycle stand up by itself?\n",
      "\n",
      "Because it was two-tired.\n",
      "\n",
      "Why was the math book sad? Because it had too many problems.\n",
      "---\n",
      "\n",
      "Total Tokens: 0\n",
      "Prompt Tokens: 0\n",
      "Completion Tokens: 0\n",
      "Total Cost (USD): $0.0\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.callbacks import get_openai_callback\n",
    "from langchain_openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model_name=\"gpt-3.5-turbo-instruct\")\n",
    "\n",
    "with get_openai_callback() as cb:\n",
    "    for chunk in llm.stream(\"Tell me a joke\"):\n",
    "        print(chunk, end=\"\", flush=True)\n",
    "    print(result)\n",
    "    print(\"---\")\n",
    "print()\n",
    "\n",
    "print(f\"Total Tokens: {cb.total_tokens}\")\n",
    "print(f\"Prompt Tokens: {cb.prompt_tokens}\")\n",
    "print(f\"Completion Tokens: {cb.completion_tokens}\")\n",
    "print(f\"Total Cost (USD): ${cb.total_cost}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0_1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
